package org.apache.spark.ml.timeseries

import org.apache.commons.math3.random.MersenneTwister
import org.apache.spark.SparkFunSuite
import org.apache.spark.ml.linalg.{DenseVector, Matrices, Vectors}
import org.apache.spark.ml.util.DefaultReadWriteTest
import org.apache.spark.mllib.util.MLlibTestSparkContext
import org.apache.spark.mllib.util.TestingUtils._


/**
  * Created by endy on 16-12-21.
  */
class UnivariateTimeSeriesSuite extends SparkFunSuite with MLlibTestSparkContext
  with DefaultReadWriteTest {

  test("lagIncludeOriginalsTrue") {
    val lagMatrix = UnivariateTimeSeries.lag(Vectors.dense(1.0, 2.0, 3.0, 4.0, 5.0), 2, true)
    assert(lagMatrix === Matrices.dense(3, 3, Array(3.0, 4.0, 5.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0)))
  }

  test("lagIncludeOriginalsFalse") {
    val lagMatrix = UnivariateTimeSeries.lag(Vectors.dense(1.0, 2.0, 3.0, 4.0, 5.0), 2, false)
    assert(lagMatrix == Matrices.dense(3, 2, Array(2.0, 3.0, 4.0, 1.0, 2.0, 3.0)))
  }

  test("autocorr") {
    val rand = new MersenneTwister(5L)
    val iidAutocorr = UnivariateTimeSeries.autocorr(Array.fill(10000)(rand.nextDouble * 5.0), 3)
    iidAutocorr.foreach(x => assert(math.abs(x) < .03))
  }

  test("upsampling") {
    // replicating upsampling examples
    // from http://www.mathworks.com/help/signal/ref/upsample.html?searchHighlight=upsample
    val y = new DenseVector(Array(1.0, 2.0, 3.0, 4.0))
    val yUp1 = UnivariateTimeSeries.upsample(y, 3, useZero = true).toArray
    assert(yUp1 === Array(1.0, 0.0, 0.0, 2.0, 0.0, 0.0, 3.0, 0.0, 0.0, 4.0, 0.0, 0.0))

    val yUp2 = UnivariateTimeSeries.upsample(y, 3, useZero = true, phase = 2).toArray
    assert(yUp2 === Array(0.0, 0.0, 1.0, 0.0, 0.0, 2.0, 0.0, 0.0, 3.0, 0.0, 0.0, 4.0))
  }

  test("downsampling") {
    // replicating downsampling examples
    // from http://www.mathworks.com/help/signal/ref/downsample.html?searchHighlight=downsample
    val y = new DenseVector((1 to 10).toArray.map(_.toDouble))
    val yDown1 = UnivariateTimeSeries.downsample(y, 3).toArray
    assert(yDown1 === Array(1.0, 4.0, 7.0, 10.0))

    val yDown2 = UnivariateTimeSeries.downsample(y, 3, phase = 2).toArray
    assert(yDown2 === Array(3.0, 6.0, 9.0))
  }

  test("signal reconstruction with spline") {
    // If we have a frequent signal, downsample it (at a rate that doesn't cause aliasing)
    // and we upsample, and apply a filter (interpolation), then the result should be fairly
    // close to the original signal. In our case, we drop NAs that are not filled by interpolation
    // (i.e no extrapolation)

    val y = (1 to 1000).toArray.map(_.toDouble / 100.0).map(Math.sin)
    val vy = new DenseVector(y)
    val lessFreq = UnivariateTimeSeries.downsample(vy, 100)
    val moreFreq = UnivariateTimeSeries.upsample(lessFreq, 100)

    // work on copies
    val splineY = UnivariateTimeSeries.fillSpline(new DenseVector(moreFreq.toArray)).toArray
    val lineY = UnivariateTimeSeries.fillLinear(new DenseVector(moreFreq.toArray)).toArray

    val MSE = (est: Array[Double], obs: Array[Double]) => {
      val errs = est.zip(obs).filter(!_._1.isNaN).map { case (yhat, yi) =>
        (yhat - yi) * (yhat - yi)
      }
      errs.sum / errs.length
    }

    val sE = MSE(splineY, y)
    val lE = MSE(lineY, y)

    // a cubic spline should be better than linear interpolation
    assert(sE < lE)
  }

  test("differencing at lag") {
    val rand = new MersenneTwister(10L)
    val n = 100
    val sampled = new DenseVector(Array.fill(n)(rand.nextGaussian))
    val lag = 5
    val diffed = UnivariateTimeSeries.differencesAtLag(sampled, lag)
    val invDiffed = UnivariateTimeSeries.inverseDifferencesAtLag(diffed, lag)

    for (i <- 0 until n) {
      assert(sampled(i) ~== invDiffed(i) absTol 1e-6)
    }

    assert(diffed(10) == (sampled(10) - sampled(5)))
    assert(diffed(99) == (sampled(99) - sampled(94)))
  }

  test("differencing of order d") {
    val rand = new MersenneTwister(10L)
    val n = 100
    val sampled = new DenseVector(Array.fill(n)(rand.nextGaussian))
    // differencing at order 1 and lag 1 should be the same
    val diffedOfOrder1 = UnivariateTimeSeries.differencesOfOrderD(sampled, 1)
    val diffedAtLag1 = UnivariateTimeSeries.differencesAtLag(sampled, 1)

    for (i <- 0 until n) {
      assert(diffedAtLag1(i) ~== diffedOfOrder1(i) absTol 1e-6)
    }

    // differencing at order and inversing should return the original series
    val diffedOfOrder5 = UnivariateTimeSeries.differencesOfOrderD(sampled, 5)
    val invDiffedOfOrder5 = UnivariateTimeSeries.inverseDifferencesOfOrderD(diffedOfOrder5, 5)

    for (i <- 0 until n) {
      assert(invDiffedOfOrder5(i) ~== sampled(i) absTol 1e-6)
    }

    // Differencing of order n + 1 should be the same as differencing one time a
    // vector that has already been differenced to order n
    val diffedOfOrder6 = UnivariateTimeSeries.differencesOfOrderD(sampled, 6)
    val diffedOneMore = UnivariateTimeSeries.differencesOfOrderD(diffedOfOrder5, 1)
    // compare start at index = 6
    for (i <- 6 until n) {
      assert(diffedOfOrder6(i) ~== diffedOneMore(i) absTol 1e-6)
    }
  }
}
